import os, sys 
CurrDir = os.path.dirname(__file__)
sys.path.append(os.path.dirname(CurrDir))

from weiqi.data.paraller_processor import GoDataProcessor
from weiqi.dlgo.encoder import OnePlaneEncoder
from weiqi.dlgo.networks.small import Layers
from keras.models import Sequential
from keras.layers.core import Dense
from keras.callbacks import ModelCheckpoint

def run():
    rows, cols = 19,19
    num_classes = rows * cols 
    num_games = 100
    encoder = OnePlaneEncoder((rows,cols))
    processor = GoDataProcessor(encoder_name=encoder.name)
    train_generator = processor.LoadData('train', num_samples=num_games, use_generator=True)
    test_generator =processor.LoadData('test', num_games, True)
    input_shape = (encoder.num_planes,rows, cols)
    network_layers = Layers(input_shape=input_shape)
    model = Sequential()
    for layer in network_layers:
        model.add(layer)
    model.add(Dense(num_classes, activation='softmax'))
    model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
    epochs = 2
    batch_size = 128
    print('Start fit_generator:')
    model.fit_generator(generator=train_generator.Generate(batch_size, num_classes),
                        epochs=epochs, steps_per_epoch=train_generator.GetNumberSamples() / batch_size,
                        validation_data=test_generator.GetNumberSamples() / batch_size)
                        # callbacks=[ModelCheckpoint('../checkpoints/small_epoch_{epoch}.h5')])
    model.evaluate_generator(generator=test_generator.Generate(batch_size, num_classes),
                             steps=test_generator.GetNumberSamples()/batch_size)
    print('End!')
    

if __name__ == '__main__':
    run()